/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"

using namespace std;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
    :rootPid((PageId &) *buffer), treeHeight((int &) *(buffer + sizePageId2))
{
    // initialize index data
    rootPid = -1;
    treeHeight = 0;
    pfMode = 'r';
}

/*
 * BTreeIndex destructor
 */
BTreeIndex::~BTreeIndex()
{
    close();
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
    RC rc = pf.open(indexname, mode);
    if (rc != 0) {
        return rc;
    }
    
    // if endpid returns <= 1 and mode == 'w' write buffer out to page 0 and initialize the root leaf, else if not empty, read buffer in, else err out
    if (mode == 'w' && pf.endPid() <= 1) {
        rootPid = 1;
        treeHeight = 1;
        
        rc = pf.write(0, (void *) buffer);
        if (rc != 0) {
            return rc;
        }
        
        BTLeafNode root;
        
        rc = root.write(1, pf);
        if (rc != 0) {
            return rc;
        }
    } 
    else if (1 < pf.endPid()) {
        rc = pf.read(0, (void *) buffer);
        if (rc != 0) {
            return rc;
        }
    }
    else {
        return RC_INVALID_FILE_FORMAT;
    }
    
    // set mode
    pfMode = mode;
    
    return 0;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
    RC rc;
    // write out index info to page 0 if opened in write mode
    if (pfMode == 'w') {
        rc = pf.write(0, (void *) buffer);
        if (rc != 0) {
            return rc;
        }
    }
    
    // close pagefile, return error if failure
    rc = pf.close();
    if (rc != 0) {
        return rc;
    }
    
    // reset mode to 'r'
    pfMode = 'r';
    
    return 0;
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
    // if not opened in write mode, err out
    if (pfMode != 'w' && pfMode != 'W') {
        return RC_INVALID_FILE_MODE;
    }

    // If root/height are at default values (-1, 0) err out
    if (rootPid == -1 || treeHeight == 0){
        return RC_INDEX_NOT_INITIALIZED;
    }
    
    RC rc;
    int midKey;
    PageId newNode;
    
    // call insertRecursive with height 1 and the appropriate node
    rc = insertRecursive(rootPid, key, rid, 1, midKey, newNode);
    // If node was split, create a new root, initialize with current root, midkey, and newNode, write it out, and update root pointer and height
    if (rc != 0) {
        if (rc == RC_NODE_SPLIT) {
            BTNonLeafNode newRoot;
            PageId newRootPid = pf.endPid();
            
            newRoot.initializeRoot(rootPid, midKey, newNode);
            rc = newRoot.write(newRootPid, pf);
            if (rc != 0) {
                return rc;
            }
            
            rootPid = newRootPid;
            treeHeight++;
        }
        else {
            return rc;
        }
    }

    return 0;
}

RC BTreeIndex::insertRecursive(PageId node, int key, const RecordId& rid, int curHeight, int& midKey, PageId& newNode){
    // if node is 0, err out
    if (node == 0) {
        return RC_INVALID_PID;
    }
    
    int rc;

    if (curHeight == treeHeight){
        BTLeafNode lnode1;
        
        rc = lnode1.read(node, pf);
        if (rc != 0) {
            return rc;
        }
        
        // Attempt to insert the node
        rc = lnode1.insert(key, rid);
        if (rc != 0){
            // If it fails because node is full, attempt insert and split, else, err out
            if (rc == RC_NODE_FULL){
                BTLeafNode lnode2;
                newNode = pf.endPid();
                
                // Attempt insert and split
                rc = lnode1.insertAndSplit(key, rid, lnode2, midKey);
                if (rc != 0){
                    return rc;
                }
                
                // update new node to point to old node's next node, and write it out
                lnode2.setNextNodePtr(lnode1.getNextNodePtr());
                rc = lnode2.write(newNode, pf);
                if (rc != 0){
                    return rc;
                }
                
                // update old node to point to new node, and write it out
                lnode1.setNextNodePtr(newNode);
                rc = lnode1.write(node, pf);
                if (rc != 0){
                    return rc;
                }
                
                // Return indicates a split
                return RC_NODE_SPLIT;
            }
            else{
                return rc;
            }
        }
        
        // write out updated node
        rc = lnode1.write(node, pf);
        if (rc != 0){
            return rc;
        }
    }
    else {
        BTNonLeafNode nlnode1;
        PageId child;
        PageId childNewNode;
        int childMidKey;
        
        rc = nlnode1.read(node, pf);
        if (rc != 0) {
            return rc;
        }
        
        // locate the appropriate child, and recurse with updated height
        nlnode1.locateChildPtr(key, child);
        rc = insertRecursive(child, key, rid, curHeight + 1, childMidKey, childNewNode);
        
        if (rc != 0){
            // If the child split, attempt to insert, else err out
            if (rc == RC_NODE_SPLIT){
                rc = nlnode1.insert(childMidKey, childNewNode);
                if(rc != 0){
                    // If the node was full,
                    if (rc == RC_NODE_FULL){
                        BTNonLeafNode nlnode2;
                        newNode = pf.endPid();
                        
                        // Attempt insert and split
                        rc = nlnode1.insertAndSplit(childMidKey, childNewNode, nlnode2, midKey);
                        if (rc != 0){
                            return rc;
                        }
                        
                        // write out new node
                        rc = nlnode2.write(newNode, pf);
                        if (rc != 0){
                            return rc;
                        }
                        
                        // write out old node
                        rc = nlnode1.write(node, pf);
                        if (rc != 0){
                            return rc;
                        }
                        
                        // Return indicates a split
                        return RC_NODE_SPLIT;
                    }
                    else{
                        return rc;
                    }
                }
                // write out updated node
                rc = nlnode1.write(node, pf);
                if (rc != 0){
                    return rc;
                }
            }
            else{
                return rc;
            }
        }
    }
    
    return 0;
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
    // If root/height are at default values (-1, 0) err out
    if (rootPid == -1 || treeHeight == 0){
        cursor.pid = 0;
        cursor.eid = 0;
        return RC_INDEX_NOT_INITIALIZED;
    }
    
    // for all levels except the leaf level, locate the correct node and follow
    cursor.pid = rootPid;
    RC rc;
    for (int curHeight = 1; curHeight < treeHeight; curHeight++){
        BTNonLeafNode nlNode;
        rc = nlNode.read(cursor.pid, pf);
        if (rc != 0) {
            return rc;
        }
        
        nlNode.locateChildPtr(searchKey, cursor.pid); 
    }
    
    // at the leaf node, locate the correct slot, and set the index as appropriate
    BTLeafNode lNode;
    rc = lNode.read(cursor.pid, pf);
    if (rc != 0) {
        return rc;
    }
    
    rc = lNode.locate(searchKey, cursor.eid);
    if (rc != 0) {
        // if there is no record >= to the search key, point to the first record in the next leaf node
        if (rc == RC_NO_SUCH_RECORD){
            cursor.pid = lNode.getNextNodePtr();
            cursor.eid = 0;
        }
        else {
            return rc;
        }
    }
    
    return 0;
}

RC BTreeIndex::getLeftmost(IndexCursor& cursor) {
    // If root/height are at default values (-1, 0) err out
    if (rootPid == -1 || treeHeight == 0){
        cursor.pid = 0;
        cursor.eid = 0;
        return RC_INDEX_NOT_INITIALIZED;
    }
    
    // for all levels except the leaf level, follow the leftmost node pointer
    cursor.pid = rootPid;
    RC rc;
    for (int curHeight = 1; curHeight < treeHeight; curHeight++){
        BTNonLeafNode nlNode;
        rc = nlNode.read(cursor.pid, pf);
        if (rc != 0) {
            return rc;
        }
        
        cursor.pid = nlNode.getLeftmostPtr(); 
    }
    
    // at the leaf node, set cursor to point to the first element in the node
    cursor.eid = 0;
    
    return 0;

}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
    // If asking for a pid past the end of the file, err out
    if (cursor.pid >= pf.endPid()){
        return RC_INVALID_PID;
    }
    
    // if asking for pid 0, assume that this is due to the end of the tree being reached from a previous call, and err out
    if (cursor.pid == 0){
        return RC_END_OF_TREE;
    }

    // Access node
    BTLeafNode node;
    RC rc;
    
    // Access node and read entry
    rc = node.read(cursor.pid, pf);
    if (rc != 0) {
        return rc;
    }
    
    node.readEntry(cursor.eid, key, rid);
    
    // if this is not the last entry in the node, increment eid, else set the cursor to point to slot 0 in the next node
    if( cursor.eid < node.getKeyCount() - 1){
        cursor.eid++;
    }
    else {
        cursor.pid = node.getNextNodePtr();
        cursor.eid = 0;
    }
  
    return 0;
}
